import faiss
import numpy as np

# 直接创建索引
# 导入 Faiss 库并创建基础索引，如IndexFlatL2（基于欧氏距离）和IndexFlatIP（基于点积相似度）
def create_index():
    # 创建一个256维的L2索引
    index_l2 = faiss.IndexFlatL2(256)  # 使用欧氏距离
    # 创建一个256维的IP索引
    index_ip = faiss.IndexFlatIP(256)  # 使用点积相似度
    return index_l2, index_ip

# 工厂方法创建索引
# 使用faiss.index_factory创建更复杂的索引配置，可指定维度、索引类型和度量标准
def create_factory_index(dimension=256):
    # 使用工厂方法创建L2索引
    index_l2 = faiss.index_factory(dimension, "Flat", faiss.METRIC_L2)
    # 使用工厂方法创建IP索引
    index_ip = faiss.index_factory(dimension, "Flat", faiss.METRIC_INNER_PRODUCT)
    return index_l2, index_ip

# 添加向量
# 建索引后，使用index.add方法添加向量，向量会被分配默认从 0 开始的唯一 ID，并可打印索引中向量总数。
def add_vectors(index):
    vectors = np.random.rand(10000, 256)  # 生成随机向量
    index.add(vectors)  # 添加向量到索引
    print(f"Total number of vectors in the index: {index.ntotal}")

# 搜索向量
# 通过index.search方法指定查询向量和返回的最近邻数量进行搜索，返回距离D和索引I。
def search_vectors(index):
    query = np.random.rand(2, 256)  # 生成两个查询向量
    k = 2  # 返回最相似的前2个向量
    D, I = index.search(query, k)  # D为距离，I为索引
    print("Distances:", D)
    print("Indices:", I)

# 删除向量
# 利用index.remove_ids指定 ID 删除向量，或使用index.reset重置索引以删除所有向量，并打印剩余向量数量。
def remove_vectors(index):
    index.remove_ids(np.array([1, 2, 3]))  # 删除指定ID的向量
    print(f"Total number of vectors after removal: {index.ntotal}")

    index.reset()  # 删除全部向量
    print(f"Total number of vectors after reset: {index.ntotal}")

# 存储与加载索引
# 使用faiss.write_index将索引保存到文件，再用faiss.read_index从文件重新加载索引。
def save_and_load_index(index, filename="vectors.faiss"):
    faiss.write_index(index, filename)  # 保存索引到文件
    loaded_index = faiss.read_index(filename)  # 从文件加载索引
    return loaded_index